from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from biz.core.ai.tools import road_function, pole_function, ask_for_more_info
from langchain.callbacks.base import BaseCallbackHandler
from typing import List, Optional, Dict, Any
from biz.core.command.chain import get_session_history
from biz.config.settings import llm
from biz.utils.logger import logger
#这个文件是函数自动调用配置

class ToolResultCallback(BaseCallbackHandler):
    """捕获工具执行结果的回调处理器"""
    
    def __init__(self):
        self.tool_results = []
        self.current_tool_name = None
        self.current_tool_input = None
    
    def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs) -> None:
        """工具开始执行时调用"""
        self.current_tool_name = serialized.get('name', 'unknown')
        self.current_tool_input = input_str
        logger.info(f"工具开始执行: {self.current_tool_name}, 输入: {input_str}")
    
    def on_tool_end(self, output: str, **kwargs) -> None:
        """工具执行结束时调用"""
        tool_result = {
            'tool_name': self.current_tool_name,
            'input': self.current_tool_input,
            'output': output,
            'timestamp': kwargs.get('timestamp')
        }
        self.tool_results.append(tool_result)
        logger.info(f"工具执行完成: {self.current_tool_name}, 输出: {output}")
    
    def on_tool_error(self, error: Exception, **kwargs) -> None:
        """工具执行出错时调用"""
        tool_result = {
            'tool_name': self.current_tool_name,
            'input': self.current_tool_input,
            'error': str(error),
            'timestamp': kwargs.get('timestamp')
        }
        self.tool_results.append(tool_result)
        logger.error(f"工具执行出错: {self.current_tool_name}, 错误: {error}")
    
    def get_results(self) -> List[Dict[str, Any]]:
        """获取所有工具执行结果"""
        return self.tool_results
    
    def clear_results(self):
        """清空结果"""
        self.tool_results = []


prompt = ChatPromptTemplate.from_messages(
    [
        ("system", """你是一个开关灯助手，只能回答开关灯相关的问题。如果用户提出的问题与开关灯无关，需要引导用户提问开关灯相关的问题。
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    不要有任何说明、解释、注释
    请尽可能从用户输入的问题中找到合适的函数，如果未找到，请使用ask_for_more_info函数来回复用户！！
    
    使用以下格式回答问题:
    Thought: 你应该经常考虑该怎么做
    Action: 调用的函数名称必须从列表中选择一个
    Action Input: 操作的输入，不要有任何解释和说明，直接输入，不能包含注释，不能包含注释，不能包含注释，不能包含注释，不能包含注释，不能有"//"
    Final Answer: 不要任何解释，直接回复最终答案。
"""),
        # First put the history
        ("placeholder", "{chat_history}"),
        # Then the new input
        ("human", "{input}"),
        # Finally the scratchpad
        ("placeholder", "{agent_scratchpad}"),
    ]
)

tools = [road_function, pole_function, ask_for_more_info]
agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, 
    max_iterations=1,
    verbose=True)


# 创建回调处理器
tool_callback = ToolResultCallback()
# 为 agent_executor 配置回调
from biz.core.command.chain import agent_executor
agent_with_callbacks = agent_executor.with_config({"callbacks": [tool_callback]})
# 创建带回调的 runnable_with_history
from langchain_core.runnables import RunnableWithMessageHistory
runnable_with_callbacks = RunnableWithMessageHistory(
    agent_with_callbacks,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

# 获取工具执行结果
tool_results = tool_callback.get_results()
content = tool_results[0]['output'][1]
logger.info("\033[1;32;40m" + content + "\033[0m")